{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddc557b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import wfdb\n",
    "import ast\n",
    "\n",
    "def load_raw_data(df, sampling_rate, path):\n",
    "    if sampling_rate == 100:\n",
    "        data = [wfdb.rdsamp(path+f) for f in df.filename_lr]\n",
    "    else:\n",
    "        data = [wfdb.rdsamp(path+f) for f in df.filename_hr]\n",
    "    data = np.array([signal for signal, meta in data])\n",
    "    return data\n",
    "\n",
    "path = 'G:/Doctra_ECG/Dataset/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.3/'\n",
    "sampling_rate=100\n",
    "calssificatin_type = \"superclasses\"    #{\"binary\",\"superclasses\",\"subclasses\"}\n",
    "\n",
    "lead_types={\"lead-I\":[1,2,3,4,5,6,7,8,9,10,11], \"bipolar-limb\":[3,4,5,6,7,8,9,10,11] , \"unipolar-limb\":[0,1,2,6,7,8,9,10,11], \"limb-leads\":[6,7,8,9,10,11] , \"precordial-leads\":[0,1,2,3,4,5],\"all-lead\":[]}\n",
    "lead_name=\"all-lead\"\n",
    "\n",
    "# load and convert annotation data\n",
    "Y = pd.read_csv(path+'ptbxl_database.csv', index_col='ecg_id')\n",
    "Y.scp_codes = Y.scp_codes.apply(lambda x: ast.literal_eval(x))\n",
    "\n",
    "#\n",
    "# Load raw signal data\n",
    "X = load_raw_data(Y, sampling_rate, path)\n",
    "\n",
    "# Load scp_statements.csv for diagnostic aggregation\n",
    "agg_df = pd.read_csv(path+'scp_statements.csv', index_col=0)\n",
    "agg_df = agg_df[agg_df.diagnostic == 1]\n",
    "\n",
    "#\n",
    "def aggregate_superclass_diagnostic(y_dic):\n",
    "    tmp = []\n",
    "    for key in y_dic.keys():\n",
    "        if key in agg_df.index:\n",
    "            tmp.append(agg_df.loc[key].diagnostic_class)\n",
    "    return list(set(tmp))\n",
    "\n",
    "\n",
    "def aggregate_subclass_diagnostic(y_dic):\n",
    "    tmp = []\n",
    "    for key in y_dic.keys():\n",
    "        if key in agg_df.index:\n",
    "            tmp.append(agg_df.loc[key].diagnostic_subclass)\n",
    "    ret = list(set(tmp))\n",
    "    return ret\n",
    "\n",
    "\n",
    "if  calssificatin_type==\"superclasses\":\n",
    "    Y['diagnostic_superclass'] = Y.scp_codes.apply(aggregate_subclass_diagnostic)\n",
    "else:\n",
    "    Y['diagnostic_superclass'] = Y.scp_codes.apply(aggregate_superclass_diagnostic)\n",
    "\n",
    "\n",
    "\n",
    "# Split data into train and test\n",
    "test_fold = 10\n",
    "# Train\n",
    "X_train = X[np.where(Y.strat_fold != test_fold)]\n",
    "y_train = Y[(Y.strat_fold != test_fold)].diagnostic_superclass\n",
    "\n",
    "\n",
    "\n",
    "# Test\n",
    "X_test = X[np.where(Y.strat_fold == test_fold)]\n",
    "y_test = Y[Y.strat_fold == test_fold].diagnostic_superclass\n",
    "\n",
    "y_train= y_train.tolist()\n",
    "y_test= y_test.tolist()\n",
    "\n",
    "if  calssificatin_type==\"binary\":\n",
    "    count =0\n",
    "    for i in y_train:\n",
    "        if 'MI' in i or 'HYP' in i or 'CD' in i or 'STTC' in i:\n",
    "            y_train[count] = 0\n",
    "        elif 'NORM' in i:\n",
    "            y_train[count] = 1\n",
    "\n",
    "        else:\n",
    "            y_train[count] = 0\n",
    "\n",
    "\n",
    "        count = count+1\n",
    "\n",
    "\n",
    "    count =0\n",
    "    for i in y_test:\n",
    "        if 'MI' in i or 'HYP' in i or 'CD' in i or 'STTC' in i:\n",
    "            y_test[count] = 0\n",
    "        elif 'NORM' in i:\n",
    "            y_test[count] = 1\n",
    "        else:\n",
    "            y_test[count] = 0\n",
    "        count = count + 1\n",
    "  \n",
    "\n",
    "print(\"done\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b43971d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "def mkdir_ifnotexists(dir):\n",
    "    if os.path.exists(dir):\n",
    "        return\n",
    "    os.mkdir(dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c222c806",
   "metadata": {},
   "outputs": [],
   "source": [
    "# leads I \n",
    "list_train= []\n",
    "list_test = []\n",
    "for i in X_train:\n",
    "    t=np.delete(i,lead_types[lead_name], 1)\n",
    "    list_train.append(t)\n",
    "    \n",
    "for i in X_test:\n",
    "    t=np.delete(i,lead_types[lead_name], 1)\n",
    "    list_test.append(t)\n",
    "\n",
    "\n",
    "np.save('x_train.npy',np.array(list_train))\n",
    "np.save(+'x_test.npy',np.array(list_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ea0fd63",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
